#!/usr/bin/env python
from argparse import ArgumentParser
from iptools.ipv4 import ip2long
from utils import get_mysql_conf
import pymysql
import csv


def main():
    host, port, username, password, db, table_name = get_mysql_conf()
    parser = ArgumentParser(version='1.0')
    parser.add_argument("input", help='set input file')
    parser.add_argument('-b', '--bulk', help='set bulk size, default 10000', type=int, default=10000)
    parser.add_argument('-H', '--host', help='set MySQL host, default [%s]' % host, default=host)
    parser.add_argument('-P', '--port', help='set MySQL port, default [%s]' % port, default=port, type=int)
    parser.add_argument('-u', '--username', help='set username, default [%s]' % username, default=username)
    parser.add_argument('-p', '--password', help='set password, default [%s]' % password, default=password)
    parser.add_argument('-d', '--db', help='set db name, default [%s]' % db, default=db)
    parser.add_argument('-t', '--table', help='set table name, default [%s]' % table_name, default=table_name)
    args = parser.parse_args()
    host = args.host
    port = args.port
    username = args.username
    password = args.password
    db = args.db
    table_name = args.table
    connection = pymysql.connect(host=host,
                                 port=port,
                                 user=username,
                                 password=password,
                                 db=db,
                                 charset='utf8mb4',
                                 cursorclass=pymysql.cursors.DictCursor)
    try:
        with connection.cursor() as cursor:
            bulk = []
            drop_table_sql = 'DROP TABLE IF EXISTS ' + table_name
            create_table_sql = 'CREATE TABLE `' + table_name + '` ( `ip_start` bigint NOT NULL, `ip_end` bigint  NOT NULL, `country` char(2) NOT NULL, `stateprov` varchar(80) NOT NULL, `city` varchar(80) NOT NULL);'
            insert_sql = "INSERT INTO `" + table_name + "` (`ip_start`, `ip_end`, `country`, `stateprov`, `city`)" \
                                                        " VALUES (%s, %s, %s, %s, %s)"
            create_index_sql = 'ALTER TABLE ' + table_name + ' ADD INDEX index_ip_end ( `ip_end`)'
            cursor.execute(drop_table_sql)
            cursor.execute(create_table_sql)
            cursor.execute(create_index_sql)
            connection.commit()
            for ip_start, ip_end, country, state, city in read_files(args.input):
                if ip_start is None:
                    continue
                bulk.append((ip_start, ip_end, country, state, city))
                if len(bulk) > args.bulk:
                    cursor.executemany(insert_sql, bulk)
                    connection.commit()
                    bulk = []
            if len(bulk) > 0:
                cursor.executemany(insert_sql, bulk)
                connection.commit()

    finally:
        connection.close()


def read_files(f):
    with open(f, 'rb') as input_file:
        reader = csv.reader(input_file)
        for line in reader:
            ip_start = ip2long(line[0])
            ip_end = ip2long(line[1])
            country = line[2]
            state = line[3]
            city = line[4]
            yield ip_start, ip_end, country, state, city


if __name__ == '__main__':
    main()
